import numpy as np
from collections import defaultdict


class NaiveBayesClassifier:
    def __init__(self):
        self.priors = {}
        self.likelihoods = defaultdict(dict)

    def fit(self, X, y):
        # Calculate class priors
        unique_labels, counts = np.unique(y, return_counts=True)
        total_samples = len(y)
        self.priors = {label: count / total_samples for label, count in zip(unique_labels, counts)}

        # Calculate likelihoods
        for label in unique_labels:
            label_indices = np.where(y == label)
            label_features = X[label_indices]
            feature_count_per_class = np.apply_along_axis(lambda x: np.bincount(x, minlength=np.max(X) + 1), axis=0,
                                                          arr=label_features)

            # Likelihood calculation with Laplace smoothing
            self.likelihoods[label] = (feature_count_per_class + 1) / (label_features.shape[0] + np.max(X) + 1)

    def predict(self, X):
        predictions = []
        for x in X:
            class_probabilities = {}
            for label in self.priors:
                # Initialize the posterior with the prior
                posterior = np.log(self.priors[label])
                # Add log likelihoods for each feature
                for i, feature_value in enumerate(x):
                    posterior += np.log(self.likelihoods[label][i][feature_value])

                class_probabilities[label] = posterior
            # Choose the class with the maximum posterior probability
            predictions.append(max(class_probabilities, key=class_probabilities.get))
        return predictions


# Example usage:
# Feature matrix with samples as rows and features as columns
X_train = np.array([[2, 1], [1, 0], [0, 1], [1, 1]])
# Labels
y_train = np.array([0, 1, 0, 1])

# Initialize the classifier
nb_classifier = NaiveBayesClassifier()
# Train the classifier
nb_classifier.fit(X_train, y_train)

# Test data
X_test = np.array([[2, 0], [1, 1]])

# Predict the labels for the test data
predictions = nb_classifier.predict(X_test)
print("Predictions:", predictions)
